from flask import Flask, redirect, flash, url_for, render_template, make_response
from flask_wtf import FlaskForm
from wtforms import StringField, PasswordField, SubmitField, IntegerField
from wtforms.validators import DataRequired, Length, InputRequired,NumberRange
from flask_sqlalchemy import SQLAlchemy
import os
from rsa_module import RsaEnc

app = Flask(__name__)
app.config["SQLALCHEMY_DATABASE_URI"] = "*****:******" + os.path.join(app.root_path, "***.****")
app.config["SECRET_KEY"] = "**********"
db = SQLAlchemy(app)


@app.cli.command()
def create_db():
    """
    创建数据库
    :return:
    """
    db.create_all()


class Delete(FlaskForm):  # 删除账户的表单,防止csrf攻击
    id = IntegerField("学号", validators=[DataRequired(), NumberRange(10000000, 50000000)])
    password = PasswordField("教务密码", validators=[InputRequired(), Length(3, 36)])  # 教务密码可以空格为首字符
    submit = SubmitField("删除")


class Register(FlaskForm):
    """
    用户信息登记
    """
    name = StringField("姓名", validators=[DataRequired(), Length(1, 5)])
    id = IntegerField("学号", validators=[DataRequired(), NumberRange(10000000, 50000000)])
    password = PasswordField("教务密码", validators=[InputRequired(), Length(3, 36)])
    classes = StringField("班级", validators=[DataRequired()])
    submit = SubmitField("提交")


class Student(db.Model):
    """
    ORM数据模型
    """
    name = db.Column(db.String)
    id = db.Column(db.INT, primary_key=True, index=True)
    password = db.Column(db.String, nullable=False)
    classes = db.Column(db.String)


@app.route("/index", methods=["GET", "POST"])
@app.route("/", methods=["GET", "POST"])
def index():
    register = Register()
    if register.validate_on_submit():
        resp = make_response()  # 创建返回对象
        resp.autocorrect_location_header = False  # 禁用自动设置location头
        resp.status_code = 302
        if Student.query.filter_by(id=register.id.data).count():
            flash("你已经添加账号啦", "error")
            return redirect(url_for("index"))
        rsa = RsaEnc()
        student = Student(id=register.id.data, name=register.name.data, password=rsa.rsa_enc(register.password.data),
                          classes=register.classes.data)
        session = db.session()
        session.add(student)
        session.commit()
        flash("添加成功！", "success")
        resp.headers['Location'] = url_for("index")
        return resp
    return render_template("index.html", form=register)


@app.route("/delete", methods=["GET", "POST"])
def delete_page():
    delete = Delete()

    if delete.validate_on_submit():
        del_target = Student.query.filter_by(id=delete.id.data).first()
        resp = make_response()  # 创建返回对象
        resp.autocorrect_location_header = False  # 禁用自动设置location头
        resp.status_code = 302
        if del_target:
            rsa = RsaEnc()
            password_raw = rsa.rsa_dec(del_target.password)
            if password_raw == delete.password.data:
                db.session.delete(del_target)
                db.session.commit()
                flash("你的账户信息已经彻底清除。", "success")
                resp.headers['Location'] = url_for("index")
                return resp
            else:
                flash("你输入的密码错误", "warning")
                resp.headers['Location'] = url_for("delete_page")
                return resp
        else:
            flash("你输入的学号不存在", "warning")
            resp.headers['Location'] = url_for("delete_page")
            return resp
    return render_template("delete.html", form=delete)


if __name__ == '__main__':
    app.run()
